// main.cpp
#include <vulkan/vulkan.h>
#include <vulkan/vulkan.hpp>
#include <vulkan/vulkan_core.h>
#include <iostream>
#include <vector>
#include <fstream>
#include <cstring>
#include <chrono>

using namespace std;
using namespace vk;
// const int MATRIX_SIZE = 256;
#define VK_LOG_DEBUG(msg) std::cerr << msg << std::endl

// 矩阵数据结构
struct Matrix {
    std::vector<float> data;
    VkBuffer buffer;
    VkDeviceMemory memory;
};

// 初始化Vulkan实例和设备
static void initVulkan(VkInstance& instance, VkDevice& device, VkPhysicalDevice& physicalDevice, VkQueue& queue) {
    // 创建实例（简化版，未包含验证层）
    VkApplicationInfo appInfo{};
    appInfo.sType = VK_STRUCTURE_TYPE_APPLICATION_INFO;
    appInfo.pApplicationName = "VulkanMatMul";
    appInfo.applicationVersion = VK_MAKE_VERSION(1, 0, 0);
    appInfo.pEngineName = "No Engine";
    appInfo.engineVersion = VK_MAKE_VERSION(1, 0, 0);
    appInfo.apiVersion = VK_API_VERSION_1_0;


    VkInstanceCreateInfo createInfo{};
    createInfo.sType = VK_STRUCTURE_TYPE_INSTANCE_CREATE_INFO;
    createInfo.pApplicationInfo = &appInfo;
    //const char* layers[] = {"VK_LAYER_KHRONOS_validation"};
    //createInfo.enabledLayerCount = 1;
    //createInfo.ppEnabledLayerNames = layers;

    vkCreateInstance(&createInfo, nullptr, &instance);

    // 选择物理设备
    uint32_t deviceCount = 0;
    vkEnumeratePhysicalDevices(instance, &deviceCount, nullptr);
    std::vector<VkPhysicalDevice> devices(deviceCount);
    vkEnumeratePhysicalDevices(instance, &deviceCount, devices.data());
    physicalDevice = devices[0]; // 默认选第一个

    // 获取物理设备信息
    for (const auto& device : devices) {
        VkPhysicalDeviceProperties props;
        vkGetPhysicalDeviceProperties(device, &props);
        std::cout << "Device Name: " << props.deviceName 
                << ", Type: " << props.deviceType 
                << ", Vulkan Version: " 
                << VK_VERSION_MAJOR(props.apiVersion) << "."
                << VK_VERSION_MINOR(props.apiVersion) << std::endl;
    }

    // 创建逻辑设备（启用计算队列）
    uint32_t queueFamilyIndex = 0;
    vkGetPhysicalDeviceQueueFamilyProperties(physicalDevice, &queueFamilyIndex, nullptr);
    std::vector<VkQueueFamilyProperties> queueFamilies(queueFamilyIndex);
    vkGetPhysicalDeviceQueueFamilyProperties(physicalDevice, &queueFamilyIndex, queueFamilies.data());

    for (uint32_t i = 0; i < queueFamilyIndex; i++) {
        if (queueFamilies[i].queueFlags & VK_QUEUE_COMPUTE_BIT) {
            queueFamilyIndex = i;
            break;
        }
    }

    float queuePriority = 1.0f;
    VkDeviceQueueCreateInfo queueCreateInfo{};
    queueCreateInfo.sType = VK_STRUCTURE_TYPE_DEVICE_QUEUE_CREATE_INFO;
    queueCreateInfo.queueFamilyIndex = queueFamilyIndex;
    queueCreateInfo.queueCount = 1;
    queueCreateInfo.pQueuePriorities = &queuePriority;

    VkDeviceCreateInfo deviceCreateInfo{};
    deviceCreateInfo.sType = VK_STRUCTURE_TYPE_DEVICE_CREATE_INFO;
    deviceCreateInfo.queueCreateInfoCount = 1;
    deviceCreateInfo.pQueueCreateInfos = &queueCreateInfo;

    try {
        vkCreateDevice(physicalDevice, &deviceCreateInfo, nullptr, &device);
    } catch (const vk::SystemError& e) {
        std::cerr<<e.what()<<std::endl;
    }
   

    VK_LOG_DEBUG("after vkGetDeviceQueue" << queue <<"\n");
    try {
        vkGetDeviceQueue(device, queueFamilyIndex, 0, &queue);
    }  catch (const vk::SystemError& e) {
        std::cerr<<e.what()<<std::endl;
    }
    VK_LOG_DEBUG("after vkGetDeviceQueue\n");
}

// 创建存储缓冲区
static void createBuffer(VkDevice device, VkPhysicalDevice physicalDevice, Matrix& matrix, VkBufferUsageFlags usage) {

    VkBufferCreateInfo bufferInfo{};
    bufferInfo.sType = VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO;
    bufferInfo.size = matrix.data.size() * sizeof(float);
    bufferInfo.usage = usage;
    bufferInfo.sharingMode = VK_SHARING_MODE_EXCLUSIVE;
    vkCreateBuffer(device, &bufferInfo, nullptr, &matrix.buffer);

    VkMemoryRequirements memRequirements;
    vkGetBufferMemoryRequirements(device, matrix.buffer, &memRequirements);

    VkMemoryAllocateInfo allocInfo{};
    allocInfo.sType = VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO;
    allocInfo.allocationSize = memRequirements.size;

    // 动态选择内存类型
    VkPhysicalDeviceMemoryProperties memProperties;
    vkGetPhysicalDeviceMemoryProperties(physicalDevice, &memProperties);
    for (uint32_t i = 0; i < memProperties.memoryTypeCount; i++) {
        if ((memRequirements.memoryTypeBits & (1 << i)) &&
            (memProperties.memoryTypes[i].propertyFlags & VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT) &&
            (memProperties.memoryTypes[i].propertyFlags & VK_MEMORY_PROPERTY_HOST_COHERENT_BIT)) {
            allocInfo.memoryTypeIndex = i;
            break;
        }
    }

    vkAllocateMemory(device, &allocInfo, nullptr, &matrix.memory);
    vkBindBufferMemory(device, matrix.buffer, matrix.memory, 0);

    // 复制数据到GPU
    void* data;
    vkMapMemory(device, matrix.memory, 0, bufferInfo.size, 0, &data);
    memcpy(data, matrix.data.data(), bufferInfo.size);
    vkUnmapMemory(device, matrix.memory);
}

// 创建Fence的正确流程
static VkFence createFence(VkDevice device) {
    VkFenceCreateInfo fenceInfo{};
    fenceInfo.sType = VK_STRUCTURE_TYPE_FENCE_CREATE_INFO;
    // 默认初始状态为未触发（unsignaled）

    VkFence fence;
    if (vkCreateFence(device, &fenceInfo, nullptr, &fence) != VK_SUCCESS) {
        throw std::runtime_error("Failed to create fence!");
    }
    return fence;
}

int main() {
    


    VK_LOG_DEBUG("Hello Vulkan!\n");
    VkInstance instance;
    VkDevice device;
    VkPhysicalDevice physicalDevice;
    VkQueue queue;
    initVulkan(instance, device, physicalDevice, queue);
    VK_LOG_DEBUG("Vulkan initialized.\n");

    // // 初始化矩阵
    // Matrix a, b, c;
    // a.data = std::vector<float>(MATRIX_SIZE * MATRIX_SIZE, 1.0f);
    // b.data = std::vector<float>(MATRIX_SIZE * MATRIX_SIZE, 2.0f);
    // c.data = std::vector<float>(MATRIX_SIZE * MATRIX_SIZE, 0.0f);

    // // 创建缓冲区
    // createBuffer(device, physicalDevice, a, VK_BUFFER_USAGE_STORAGE_BUFFER_BIT);
    // createBuffer(device, physicalDevice, b, VK_BUFFER_USAGE_STORAGE_BUFFER_BIT);
    // createBuffer(device, physicalDevice, c, VK_BUFFER_USAGE_STORAGE_BUFFER_BIT | VK_BUFFER_USAGE_TRANSFER_DST_BIT);

    // VK_LOG_DEBUG("Matrix initialized.");

    // // 加载计算着色器（需预先编译为SPV）
    // VkShaderModule shaderModule;
    // std::ifstream file("matrix.comp.spv", std::ios::binary);
    // std::vector<char> code((std::istreambuf_iterator<char>(file)), std::istreambuf_iterator<char>());
    // VkShaderModuleCreateInfo shaderInfo{};
    // shaderInfo.sType = VK_STRUCTURE_TYPE_SHADER_MODULE_CREATE_INFO;
    // shaderInfo.codeSize = code.size();
    // shaderInfo.pCode = reinterpret_cast<const uint32_t*>(code.data());
    // vkCreateShaderModule(device, &shaderInfo, nullptr, &shaderModule);

    // VK_LOG_DEBUG("Shader loaded.");

    // // 创建描述符集布局
    // VkDescriptorSetLayoutBinding bindings[3];
    // for (int i = 0; i < 3; i++) {
    //     bindings[i].binding = i;
    //     bindings[i].descriptorType = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER;
    //     bindings[i].descriptorCount = 1;
    //     bindings[i].stageFlags = VK_SHADER_STAGE_COMPUTE_BIT;
    //     bindings[i].pImmutableSamplers = nullptr;
    // }

    // VK_LOG_DEBUG("Descriptor set layout created.");

    // VkDescriptorSetLayoutCreateInfo layoutInfo{};
    // layoutInfo.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_LAYOUT_CREATE_INFO;
    // layoutInfo.bindingCount = 3;
    // layoutInfo.pBindings = bindings;
    // VkDescriptorSetLayout descriptorSetLayout;
    // vkCreateDescriptorSetLayout(device, &layoutInfo, nullptr, &descriptorSetLayout);

    // VK_LOG_DEBUG("Descriptor set layout created.");

    

    // // 创建计算管线
    // VkPipelineLayoutCreateInfo pipelineLayoutInfo{};
    // pipelineLayoutInfo.sType = VK_STRUCTURE_TYPE_PIPELINE_LAYOUT_CREATE_INFO;
    // pipelineLayoutInfo.setLayoutCount = 1;
    // pipelineLayoutInfo.pSetLayouts = &descriptorSetLayout;
    // VkPipelineLayout pipelineLayout;
    // vkCreatePipelineLayout(device, &pipelineLayoutInfo, nullptr, &pipelineLayout);

    // VK_LOG_DEBUG("Pipeline layout created.");

    // VkComputePipelineCreateInfo pipelineInfo{};
    // pipelineInfo.sType = VK_STRUCTURE_TYPE_COMPUTE_PIPELINE_CREATE_INFO;
    // pipelineInfo.stage.sType = VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO;
    // pipelineInfo.stage.stage = VK_SHADER_STAGE_COMPUTE_BIT;
    // pipelineInfo.stage.module = shaderModule;
    // pipelineInfo.stage.pName = "main";
    // pipelineInfo.layout = pipelineLayout;
    // VkPipeline pipeline;
    // vkCreateComputePipelines(device, VK_NULL_HANDLE, 1, &pipelineInfo, nullptr, &pipeline);

    // VK_LOG_DEBUG("Pipeline created.");

    // // 创建描述符池和描述符集
    // VkDescriptorPoolSize poolSize{};
    // poolSize.type = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER;
    // poolSize.descriptorCount = 3;

    // VkDescriptorPoolCreateInfo poolInfo{};
    // poolInfo.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_POOL_CREATE_INFO;
    // poolInfo.poolSizeCount = 1;
    // poolInfo.pPoolSizes = &poolSize;
    // poolInfo.maxSets = 1;
    // VkDescriptorPool descriptorPool;
    // vkCreateDescriptorPool(device, &poolInfo, nullptr, &descriptorPool);

    // VkDescriptorSetAllocateInfo allocInfo{};
    // allocInfo.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_ALLOCATE_INFO;
    // allocInfo.descriptorPool = descriptorPool;
    // allocInfo.descriptorSetCount = 1;
    // allocInfo.pSetLayouts = &descriptorSetLayout;
    // VkDescriptorSet descriptorSet;
    // vkAllocateDescriptorSets(device, &allocInfo, &descriptorSet);

   

    // VkDescriptorBufferInfo bufferInfos[3] = {
    //     {a.buffer, 0, VK_WHOLE_SIZE},
    //     {b.buffer, 0, VK_WHOLE_SIZE},
    //     {c.buffer, 0, VK_WHOLE_SIZE}
    // };

    // VkWriteDescriptorSet writes[3];
    // memset(writes, 0, sizeof(writes)); // 初始化为
    // for (int i = 0; i < 3; i++) {
    //     writes[i].sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET;
    //     writes[i].pNext = nullptr; // 显式设置pNext
    //     writes[i].dstSet = descriptorSet;
    //     writes[i].dstBinding = i;
    //     writes[i].descriptorCount = 1;
    //     writes[i].descriptorType = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER;
    //     writes[i].pBufferInfo = &bufferInfos[i];
    // }
    
    // vkUpdateDescriptorSets(device, 3, writes, 0, nullptr);
    
    // VK_LOG_DEBUG("Descriptor set created.");
    
    // // 录制命令缓冲区
    // VkCommandPool commandPool;
    // VkCommandPoolCreateInfo poolCreateInfo{};
    // poolCreateInfo.sType = VK_STRUCTURE_TYPE_COMMAND_POOL_CREATE_INFO;
    // poolCreateInfo.queueFamilyIndex = 0; // 需与队列族一致
    // vkCreateCommandPool(device, &poolCreateInfo, nullptr, &commandPool);

    // VkCommandBuffer commandBuffer;
    // VkCommandBufferAllocateInfo allocCmdBufInfo{};
    // allocCmdBufInfo.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO;
    // allocCmdBufInfo.commandPool = commandPool;
    // allocCmdBufInfo.level = VK_COMMAND_BUFFER_LEVEL_PRIMARY;
    // allocCmdBufInfo.commandBufferCount = 1;
    // vkAllocateCommandBuffers(device, &allocCmdBufInfo, &commandBuffer);

    // VkCommandBufferBeginInfo beginInfo{};
    // beginInfo.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO;
    // vkBeginCommandBuffer(commandBuffer, &beginInfo);

    // vkCmdBindPipeline(commandBuffer, VK_PIPELINE_BIND_POINT_COMPUTE, pipeline);
    // vkCmdBindDescriptorSets(commandBuffer, VK_PIPELINE_BIND_POINT_COMPUTE, pipelineLayout, 0, 1, &descriptorSet, 0, nullptr);
    // vkCmdDispatch(commandBuffer, MATRIX_SIZE/16, MATRIX_SIZE/16, 1); // 工作组大小16x16 [[6]]

    // vkEndCommandBuffer(commandBuffer);

    // // 提交命令并同步
    // VkSubmitInfo submitInfo{};
    // submitInfo.sType = VK_STRUCTURE_TYPE_SUBMIT_INFO;
    // submitInfo.commandBufferCount = 1;
    // submitInfo.pCommandBuffers = &commandBuffer;

    //     // 在命令提交时使用
    // VkFence fence = createFence(device);

    // auto start = std::chrono::high_resolution_clock::now();
    // vkQueueSubmit(queue, 1, &submitInfo, fence); // 提交时关联Fence
    // vkWaitForFences(device, 1, &fence, VK_TRUE, UINT64_MAX);
    // auto end = std::chrono::high_resolution_clock::now();
    // std::chrono::duration<double> elapsed = end - start;
    // std::cout << "Total time: " << elapsed.count() * 1000 << " ms" << std::endl;

    // // 读取结果
    // void* mappedMemory;
    // vkMapMemory(device, c.memory, 0, VK_WHOLE_SIZE, 0, &mappedMemory);
    // std::vector<float> result(MATRIX_SIZE * MATRIX_SIZE);
    // memcpy(result.data(), mappedMemory, MATRIX_SIZE * MATRIX_SIZE * sizeof(float));
    // vkUnmapMemory(device, c.memory);

    // // 验证结果（预期每个元素为256*2=512）
    // std::cout << "Result[0][0] = " << result[0] << std::endl;
    // // 在memcpy获取结果后添加以下代码：
    // std::cout << "Matrix Result:\n";
    // for (int row = 0; row < MATRIX_SIZE; ++row) {
    //     for (int col = 0; col < MATRIX_SIZE; ++col) {
    //         std::cout << result[row * MATRIX_SIZE + col] << "\t";
    //     }
    //     std::cout << "\n";
    // }

    // // 清理资源（省略部分代码）
    // vkDestroyFence(device, fence, nullptr);
    // vkDestroyCommandPool(device, commandPool, nullptr);
    // vkDestroyPipeline(device, pipeline, nullptr);
    // vkDestroyPipelineLayout(device, pipelineLayout, nullptr);
    // vkDestroyDescriptorPool(device, descriptorPool, nullptr);
    // vkDestroyDescriptorSetLayout(device, descriptorSetLayout, nullptr);
    // vkDestroyShaderModule(device, shaderModule, nullptr);
    // vkDestroyBuffer(device, a.buffer, nullptr);
    // vkFreeMemory(device, a.memory, nullptr);
    // vkDestroyBuffer(device, b.buffer, nullptr);
    // vkFreeMemory(device, b.memory, nullptr);
    // vkDestroyBuffer(device, c.buffer, nullptr);
    // vkFreeMemory(device, c.memory, nullptr);
    // vkDestroyDevice(device, nullptr);
    // vkDestroyInstance(instance, nullptr);

    return 0;
}